package com.xh.common.thread.ext.queue;

import java.util.ArrayDeque;
import java.util.Comparator;
import java.util.Queue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 自定义的优先级队列，支持阻塞获取不同等级的队列。
 * 如队列优先级有5级，take(0) 可以获取0,1,2,3,4优先级的项目，take(2) 可以获取2,3,4优先级的项目。
 * 实现compare接口时必须保证level值高的排在前面。
 * 为每个级别设置了不同锁的条件。
 * 系统的队列实现不是一个稳定的排序。
 * 系统的队列并不是完整排序，会在放入元素和获取元素时动态排序部分。
 * offerWait操作是专门用于自定义线程池时将任务分配为空闲线程的，优先级最高。
 */
public class LevelPriorityBlockingQueue<E extends Level> extends PriorityBlockingQueue<E>
        implements LevelBlockingQueue<E> {

    // 队列的等级
    private int mLevel;
    // 锁条件的数组
    private Condition[] mNotEmpty;
    // 等待获取数据的线程的数量
    private int[] mWaitCount;
    // 等待队列，用于给level等待线程指定执行任务
    private Queue<E>[] mWaitQueue;

    public LevelPriorityBlockingQueue(int level) {
        super();
        init(level);
    }

    public LevelPriorityBlockingQueue(int level, int initialCapacity,
                                      Comparator<? super E> comparator) {
        super(initialCapacity, comparator);
        init(level);
    }

    /**
     * 初始化优先级和锁
     *
     * @param level
     */
    private void init(int level) {
        if (level < 0) {
            throw new IllegalArgumentException();
        }
        mLevel = level;
        mNotEmpty = new Condition[mLevel];
        for (int i = 0; i < mLevel; i++) {
            mNotEmpty[i] = lock.newCondition();
        }

        mWaitCount = new int[mLevel];
        mWaitQueue = new ArrayDeque[mLevel];
        for (int i = 0; i < mLevel; i++) {
            mWaitQueue[i] = new ArrayDeque<>(2);
        }
    }

    @Override
    public E take() throws InterruptedException {
        return take(0);
    }

    @Override
    public E take(int level) throws InterruptedException {
        checkAvailableLevel(level);

        final ReentrantLock lock = this.lock;
        lock.lockInterruptibly();
        E result;
        try {
            mWaitCount[level]++;
            while ((result = dequeue(level)) == null) {
                mNotEmpty[level].await();
            }
        } finally {
            mWaitCount[level]--;
            lock.unlock();
        }
        return result;
    }

    @Override
    public E poll() {
        return poll(0);
    }

    @Override
    public E poll(int level) {
        checkAvailableLevel(level);

        final ReentrantLock lock = this.lock;
        lock.lock();
        try {
            return dequeue(level);
        } finally {
            lock.unlock();
        }
    }

    @Override
    public E poll(long timeout, TimeUnit unit) throws InterruptedException {
        return poll(0, timeout, unit);
    }

    @Override
    public E poll(int level, long timeout, TimeUnit unit) throws InterruptedException {
        checkAvailableLevel(level);

        long nanos = unit.toNanos(timeout);
        final ReentrantLock lock = this.lock;
        lock.lockInterruptibly();
        E result;
        try {
            mWaitCount[level]++;
            while ((result = dequeue(level)) == null && nanos > 0)
                nanos = mNotEmpty[level].awaitNanos(nanos);
        } finally {
            mWaitCount[level]--;
            lock.unlock();
        }
        return result;
    }

    @Override
    public E peek(int level) {
        final ReentrantLock lock = this.lock;
        lock.lock();
        try {
            if (size == 0) {
                return null;
            }

            E item = (E) queue[0];
            if (level <= item.getLevel()) {
                return item;
            } else {
                return null;
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public E peek() {
        return peek(0);
    }

    @Override
    public boolean offer(E e) {
        if (e == null) {
            throw new NullPointerException();
        }
        int level = e.getLevel();
        checkAvailableLevel(level);

        final ReentrantLock lock = this.lock;
        lock.lock();
        int n, cap;
        Object[] array;
        while ((n = size) >= (cap = (array = queue).length)) {
            tryGrow(array, cap);
        }
        try {
            Comparator<? super E> cmp = comparator();
            if (cmp == null) {
                siftUpComparable(n, e, array);
            } else {
                siftUpUsingComparator(n, e, array, cmp);
            }
            size = n + 1;

            // 唤起比任务优先级低的所有线程，唤醒的次序从最低的优先级开始，这样尽量保证被最低优先级的线程获取到。
            for (int i = 0; i <= level; i++) {
                mNotEmpty[i].signal();
            }
        } finally {
            lock.unlock();
        }
        return true;
    }

    @Override
    public boolean offerWait(E e, int level) {
        if (e == null)
            throw new NullPointerException();
        checkAvailableLevel(level);
        checkAvailableLevel(e.getLevel());

        final ReentrantLock lock = this.lock;
        lock.lock();

        try {
            int count = mWaitCount[level];
            Queue<E> queue = mWaitQueue[level];

            // 如果存在空闲的线程，同时等待队列的任务小于空闲线程的数量时，直接把任务添加到队列中。
            // 由于线程都是空闲的，必然都会被立刻执行，所以可以不需要进行排序。
            if (count <= 0 || queue.size() >= count) {
                return false;
            }

            queue.offer(e);

            // 只需要唤醒当前等级的线程。
            mNotEmpty[level].signal();
        } finally {
            lock.unlock();
        }
        return true;
    }

    /**
     * 获取队头元素，如果等级不匹配则返回null
     *
     * @param level
     * @return
     */
    private E dequeue(int level) {
        // 先从等待队列中进行获取
        E e = mWaitQueue[level].poll();
        if (e != null) {
            return e;
        }

        if (size <= 0)
            return null;
        else {
            // 然后从正常的队列中进行获取
            if (level <= ((E) queue[0]).getLevel()) {
                return dequeue();
            }
            return null;
        }
    }

    private void checkAvailableLevel(int level) {
        if (level < 0 || level >= mLevel) {
            throw new IllegalArgumentException("Level is " + mLevel + ", request is " + level);
        }
    }
}
